import csv
import os
import warnings
from inspect import currentframe, getframeinfo

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
from multiprocessing import Process

import scipy.io as scio

import numpy as np
import torch
from DSPPytorch import *

import util

BASE_DIR = os.path.dirname(__file__)


def LE_exp(lp_range=[1]):
    frame = currentframe()
    experiment_name = getframeinfo(frame).function
    lp_range = [*lp_range]
    Qcache = np.zeros([len(lp_range)])
    pick_sym_num = 110000
    modOrder = 4
    result_save_dir = os.path.join(BASE_DIR, 'result_exp/{}'.format(experiment_name))
    if not os.path.exists(result_save_dir):
        os.makedirs(result_save_dir)

    '''将结果写入CSV文件'''
    with open(os.path.join(result_save_dir, 'Q_results_under_{}.csv'.format(experiment_name)),
              'w') as f:
        csv_writer = csv.writer(f)
        csv_writer.writerow([
            'lp', 'BER', 'Q factor', 'Q^2 factor'
        ])

        '''进入循环，对每个发射功率下的信号进行均衡'''
        for lpIndx, lp in enumerate(lp_range):
            dataPath = os.path.join(
                BASE_DIR, f'data/experiment/16QAM20GBaud1800kmHe/tstSet_lp_{lp}.mat'
            )
            '''读取文件信息'''
            data = scio.loadmat(dataPath)
            symbolRate = 20e9
            spanNum = data['spanNum'][0, 0]
            prbs = data['prbs']
            spanLen = 100e3
            L = spanLen * spanNum
            DL = -17e-6 * L
            sig = torch.from_numpy(data['sig'])
            sig = sig.reshape(1, 1, -1)

            constellations = torch.from_numpy(util.CONST_16QAM)

            '''实例化后端的DSP类'''
            # mf = MatchedFilterLayer(alpha=rolloff, span=100, sample_factor=8, case_num=2)  # span: Correlated symbols
            edc = EDCLayer(symbolRate=symbolRate, DL=DL, case_num=1)
            pr = PhaseRecLayer(1024)
            lms = FIRLayer(tap=32, case_num=1, power_norm=True, centor_one=True)

            if torch.cuda.is_available():
                sig = sig.cuda()
                lms = lms.cuda()
                edc = edc.cuda()

            '''DSP procedure for the signal through the link'''
            # sig = sig.cpu()
            # sig = mf(sig)
            # if torch.cuda.is_available():
            #     sig = sig.cuda() # 重新放回cuda上
            sig = sig[..., 0:pick_sym_num * 2]
            prbs = prbs[..., 0:pick_sym_num * modOrder]
            sig = edc(sig)

            # '''后端的TDE均衡PMD以及PDL等，是自适应的方法'''
            # sig = sig.cpu()
            # sig = FIRLayer.time_vary_infer(sig, err_mode='Godard', iter_num=2, lr=5e-4, tap=13)
            # if torch.cuda.is_available():
            #     sig = sig.cuda()  # 重新放回cuda上
            '''DownSampling and PhaseRecovery'''
            sig = pr(sig[..., 1::2])
            '''DD-LMS'''
            lms.fit(sig, err_mode='DDM', constellations=constellations, iter_num=4, block_size=4028,
                    remain=2048, lr=5e-4)
            sig = lms(sig)

            '''与Prbs进行比较计算误码率'''
            sig = sig.cpu().data.numpy().squeeze()
            sig, ber, good_rotate = util.pr_ber(sig, prbs, constellations.cpu().data.numpy())

            ber = np.mean(ber)
            Q = util.ber2q(ber)
            Q2 = util.ber2Q2(ber)
            Qcache[lpIndx] = Q

            print(f'lp {lp} done, Q factor: {Q}')

            # '''绘制星座图'''
            # p = Process(target=util.colorfulPlot, args=(sig[:, np.newaxis, ...], prbs), kwargs={'show': True})
            # p.start()

            '''保存结果'''
            csv_writer.writerow([lp, ber, Q, Q2])

            '''end'''


if __name__ == '__main__':
    warnings.filterwarnings("ignore")
    LE_exp(lp_range=[-3, -2, -1, 0, 1, 2, 3])
